/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file flash_attention_score_grad_tiling_g_bn2.h
 * \brief
 */

#pragma once

#include "tiling/tiling_base.h"
#include "flash_attention_score_grad_tiling_g_bn2_def.h"

namespace optiling {

struct BaseInfoParams { // 频繁使用的基础参数
    uint32_t coreNum;
    uint32_t aicNum;
    uint32_t ubSize;
    uint32_t l1Size;
    uint32_t l0aSize;
    uint32_t l0bSize;
    uint32_t l0cSize;

    uint32_t b;
    uint32_t n1;
    uint32_t n2;
    uint32_t s1;
    uint32_t s2;
    uint32_t g;
    uint32_t d;
    uint32_t h;
    uint32_t blockNum;
    uint32_t mask;

    uint32_t tilingKey;
    uint32_t usedUBSize;

    uint32_t layout;

    bool isSparse;
    bool isFlash;
    uint32_t dataType;
    uint32_t dataTypeSize;
    uint32_t calTypeSize;
    ge::DataType dType;

    // 基本块信息
    uint32_t gInner;
};

class FlashAttentionScoreGradUgBbn : public TilingBaseClass {
public:
  explicit FlashAttentionScoreGradUgBbn(gert::TilingContext* context_) : TilingBaseClass(context_) {
  }
  ~FlashAttentionScoreGradUgBbn() override {
  }

protected:
  bool IsCapable() override;
  ge::graphStatus GetPlatformInfo() override;
  ge::graphStatus GetShapeAttrsInfo() override;
  ge::graphStatus DoOpTiling() override;
  ge::graphStatus DoLibApiTiling() override;
  uint64_t GetTilingKey() const override;
  ge::graphStatus GetWorkspaceSize() override;
  ge::graphStatus PostTiling() override;

private:
  FlashAttentionScoreGradTilingDataUgBbn tilingData;
  BaseInfoParams mBaseInfoParams;
  bool CheckGInnerLegal(uint32_t gInner);
  uint32_t CalcGInner();

  ge::graphStatus SetBaseTilingData();
  ge::graphStatus SetCoresSplitTilingData();
  ge::graphStatus SetUbSplitTilingData();
  void SetMatmulTiling1(matmul_tiling::MatmulApiTiling& mm, TCubeTiling& mmTiling);
  void SetMatmulTiling2(matmul_tiling::MatmulApiTiling& mm, TCubeTiling& mmTiling);
  void SetMatmulTiling3(matmul_tiling::MatmulApiTiling& mm, TCubeTiling& mmTiling);
};

}  // namespace optiling
